# -*- coding: utf-8 -*-
"""
BERT相似度检索器
基于BERT模型计算问题间的语义相似度进行检索

作者: [您的姓名]
日期: 2024
"""

import os
import numpy as np
import torch
from sklearn.metrics.pairwise import cosine_similarity
from transformers import BertModel, BertTokenizer
from loguru import logger
from typing import List, Dict, Tuple
from ...config import *  # 导入配置，包括镜像设置
from ...tools.cache_manager import cache_manager

# 使用全局缓存管理器
models_cache_dir = cache_manager.cache_dir


class BertSimilarityRetriever:
    """
    基于BERT模型的相似问题检索器
    
    该检索器使用预训练的BERT模型来计算问题间的语义相似度，
    通过余弦相似度找到与查询问题最相似的候选问题。
    
    特点：
    - 支持中文BERT模型
    - 使用[CLS]标记的向量表示整个句子
    - 支持自定义相似度阈值
    - 提供详细的相似度分数
    """

    def __init__(self, model_name: str = "hfl/chinese-roberta-wwm-ext"):
        """
        初始化BERT检索器

        Args:
            model_name: BERT模型名称（从HuggingFace加载）
        """
        self.model_name = model_name
        self.tokenizer = None
        self.model = None
        self.sentences = []
        self.embeddings = None
        self.similarity_matrix = None

        # 加载模型
        self.load_model()



    def load_model(self):
        """加载BERT模型和分词器"""
        try:
            logger.info(f"🌐 从HuggingFace国内镜像加载模型: {self.model_name}")
            logger.info(f"🚀 镜像地址: {HF_MIRROR_ENDPOINT}")
            logger.info(f"📁 使用缓存目录: {models_cache_dir}")
            
            self.tokenizer = BertTokenizer.from_pretrained(
                self.model_name,
                cache_dir=models_cache_dir
            )
            self.model = BertModel.from_pretrained(
                self.model_name,
                cache_dir=models_cache_dir
            )

            self.model.eval()
            logger.info("✅ BERT模型加载成功")
        except Exception as e:
            logger.error(f"❌ 模型加载失败: {e}")
            logger.error("💡 请检查网络连接是否正常")
            self.tokenizer = None
            self.model = None

    def load_candidates(self, sentences: List[str]):
        """
        加载候选问题并计算嵌入向量

        Args:
            sentences: 候选问题列表
        """
        self.sentences = sentences

        if self.model is None:
            logger.error("❌ 模型未加载，无法计算嵌入")
            return

        logger.debug("\n生成BERT嵌入向量:")
        self.embeddings = self.get_bert_embeddings(sentences)

        logger.debug("\n计算相似度矩阵...")
        self.similarity_matrix = cosine_similarity(self.embeddings)

    def get_bert_embeddings(self, texts: List[str]) -> np.ndarray:
        """
        获取BERT嵌入向量

        Args:
            texts: 文本列表

        Returns:
            嵌入向量矩阵
        """
        # 编码文本
        encoded_input = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=128,  # 增加长度以处理长问题
            return_tensors='pt'
        )

        with torch.no_grad():
            outputs = self.model(**encoded_input)

        # 使用最后一层的[CLS]标记向量
        last_hidden_state = outputs.last_hidden_state
        cls_vectors = last_hidden_state[:, 0, :]  # 取[CLS]标记
        return cls_vectors.numpy()

    def find_similar_questions(self, user_query: str, top_k: int = 5, 
                             similarity_threshold: float = 0.0) -> List[Dict]:
        """
        查找与用户查询最相似的top_k个问题

        Args:
            user_query: 用户查询问题
            top_k: 返回相似问题的数量
            similarity_threshold: 相似度阈值，低于该值的结果将被过滤

        Returns:
            包含相似问题及其相似度分数的字典列表
        """
        if self.embeddings is None or len(self.sentences) == 0:
            logger.error("❌ 候选问题未加载")
            return []

        # 计算查询的嵌入
        query_embedding = self.get_bert_embeddings([user_query])[0]

        # 计算查询与候选的相似度
        similarities = cosine_similarity([query_embedding], self.embeddings)[0]

        # 获取最相似的top_k个问题
        top_indices = np.argsort(similarities)[::-1][:top_k]
        results = []
        for idx in top_indices:
            if similarities[idx] >= similarity_threshold:
                results.append({
                    'question': self.sentences[idx],
                    'similarity': float(similarities[idx])
                })

        return results

    def batch_similarity(self, queries: List[str], top_k: int = 5) -> Dict[str, List[Dict]]:
        """
        批量计算多个查询的相似问题

        Args:
            queries: 查询问题列表
            top_k: 每个查询返回的相似问题数量

        Returns:
            查询到相似问题的映射字典
        """
        results = {}
        for query in queries:
            results[query] = self.find_similar_questions(query, top_k=top_k)
        return results

    def get_similarity_matrix(self) -> np.ndarray:
        """返回相似度矩阵"""
        return self.similarity_matrix